from llms.zhipu import chat  # 调用国产模型的接口
from prompts.inteview_prompt import INIT_QUESTIONS_PROMPT  # 初始化问题的 Prompt
from utils.parse_utils import parse_markdown_json_content  # 解析工具
from typing_extensions import TypedDict
from typing import List, Dict, Literal
from langgraph.graph import StateGraph, START, END
from langgraph.types import interrupt, Command

MAX_FOLLOW_UPS = 2  # 最大追问次数限制


# 定义状态结构
class State(TypedDict):
    jd_detail: str  # 职位描述
    questions: List[str]  # 主问题列表
    user_answers: List[Dict]  # 用户回答列表
    graph_state: str  # 状态日志记录
    current_question: str  # 当前主问题
    current_follow_up: str  # 当前追问问题
    follow_up_count: int  # 追问次数计数器


# 初始化节点：生成问题列表
def init_interview(state: State) -> Dict[str, str]:
    print("--- 初始化阶段 ---")
    print(f"欢迎来到模拟面试！职位是：{state['jd_detail']}")

    # 使用 Prompt 模板创造问题
    prompt = INIT_QUESTIONS_PROMPT.format(jd_detail=state["jd_detail"])
    response = chat([{"role": "user", "content": prompt}])

    try:
        questions = parse_markdown_json_content(response)
        state["questions"] = [q.strip() for q in questions if q.strip()]
    except Exception as e:
        print(f"生成问题失败: {str(e)}")
        state["questions"] = []

    print("生成的问题如下：")
    for idx, q in enumerate(state["questions"], 1):
        print(f"{idx}. {q}")

    return {"graph_state": "初始化完成，已生成面试问题列表。", "questions": state["questions"],
            "follow_up_count": 0}


# 候选人节点：候选人回答问题
def candidate_response(state: State) -> Dict[str, str]:
    # 初始化需要的返回更新字段
    updated_state = {
        "current_follow_up": None,  # 默认清空追问
        "current_question": state.get("current_question"),  # 当前主问题
        "user_answers": state["user_answers"],  # 现有用户答案列表
        "questions": state["questions"],  # 剩余主问题
    }

    # 判断是回答主问题还是追问
    if state.get("current_follow_up"):
        question = state["current_follow_up"]
    elif state["questions"]:
        question = state["questions"][0]  # 获取第一个主问题的内容但不立即移除
        updated_state["current_question"] = question  # 更新当前主问题
        updated_state["follow_up_count"] = 0  # 重置追问计数器
    else:
        raise ValueError("没有问题需要回答。（错误逻辑）")

    print(f"[候选人请回答]: {question}")
    result = interrupt(
        {
            "task": f"请回答以下面试问题：{question} 要求回答简洁，回答格式{{\"answer\":\"xxxx\"}}",
            "current_question": question
        }

    )
    answer = result["answer"]
    print(f"[候选人回答]: {answer}")

    # 更新用户回答列表
    if state.get("current_follow_up"):
        # 如果是追问，则在问题前加标注“(追问)”
        updated_state["user_answers"] = state["user_answers"] + [{"question": f"(追问) {question}", "answer": answer}]
    else:
        # 如果是主问题，直接记录答案，并移除已回答的问题
        updated_state["user_answers"] = state["user_answers"] + [{"question": question, "answer": answer}]
        updated_state["questions"] = state["questions"][1:]  # 移除已回答的问题

    return {
        "graph_state": f"完成回答：{question}",
        **updated_state,  # 更新的状态内容
    }


# 条件边逻辑：决定流程下一步
def interviewer_decision_node(state: State) -> Command[Literal["candidate_response", "evaluate_interview"]]:
    # 如果所有问题已回答完毕，并且没有追问逻辑，则直接进入总结
    if not state["questions"] and not state.get("current_follow_up"):
        print("[面试官]: 所有问题和追问结束，进入总结阶段。")
        return Command(goto="evaluate_interview")

    # 如果追问次数已达到限制，直接进入下一主问题
    if state.get("follow_up_count", 0) >= MAX_FOLLOW_UPS:
        print("[面试官]: 追问次数已达上限，进入下一问题...")
        return Command(update={"current_follow_up": None}, goto="candidate_response")

    # 否则，检查是否需要追问
    latest_question = state["user_answers"][-1]["question"]
    latest_answer = state["user_answers"][-1]["answer"]

    # 调用国产语言模型生成追问
    prompt = (
        f"候选人针对以下问题回答了：'{latest_question}'。\n"
        f"回答是：'{latest_answer}'。\n请基于此生成一个后续追问问题，"
        f"如果无需追问，则返回空字符串。"
    )
    response = chat([{"role": "user", "content": prompt}])

    try:
        follow_up = response.strip()
        if follow_up:
            print(f"[面试官生成追问]: {follow_up}")
            return Command(
                update={
                    "current_follow_up": follow_up,  # 更新追问问题状态
                    "follow_up_count": state.get("follow_up_count", 0) + 1  # 增加追问次数
                },
                goto="candidate_response"  # 回到候选人响应流程
            )
    except Exception as e:
        print(f"生成追问失败: {str(e)}")

    # 如果没有追问，切回候选人节点处理下一问题
    print("[面试官]: 没有追问，进入下一问题...")
    return Command(goto="candidate_response")


# 总结节点：总结候选人的表现
def evaluate_interview(state: State) -> Dict[str, str]:
    print("--- 面试总结与反馈 ---")
    # 拼接所有回答的内容
    answers = "\n".join([f"问题: {qa['question']} - 回答: {qa['answer']}" for qa in state["user_answers"]])
    prompt = f"以下是候选人在面试中的回答内容:\n{answers}\n请总结候选人的表现并给出改进建议。"
    feedback = chat([{"role": "user", "content": prompt}])

    # 存储和输出总结
    state["graph_state"] += f"\n面试总结：\n{feedback}"
    print("[面试总结完成]:")
    print(feedback)

    return {"graph_state": state["graph_state"]}


# --------- 构建流程图 ----------
builder = StateGraph(State)

# 添加节点
builder.add_node("init_interview", init_interview)  # 初始化节点
builder.add_node("candidate_response", candidate_response)  # 候选人回答节点
builder.add_node("interviewer_decision_node", interviewer_decision_node)  # 面试官节点
builder.add_node("evaluate_interview", evaluate_interview)  # 总结节点

# 定义流程边
builder.add_edge(START, "init_interview")  # 初始化开始
builder.add_edge("init_interview", "candidate_response")  # 初始化后进入第一个问题
builder.add_edge("candidate_response", "interviewer_decision_node")  # 初始化后进入第一个问题
builder.add_edge("evaluate_interview", END)  # 总结到结束

# 编译图结构
graph = builder.compile()
